package xyz.thoughtset.viewer.common.ai.model.factory;


import org.springframework.util.MultiValueMap;
import org.springframework.util.MultiValueMapAdapter;
import xyz.thoughtset.viewer.common.ai.model.entity.AiNode;
import xyz.thoughtset.viewer.common.ai.model.entity.ModelParam;

import java.util.*;


public class ModelsRegistry {
    private static final Map<String,ModelBuilder> DEFAULT_MODEL_BUILDER_MAP = new HashMap();
    private static final MultiValueMap<String,ModelBuilder> MODEL_BUILDER_MAP = new MultiValueMapAdapter<>(new HashMap<>());

    public static void registerModel(ModelBuilder modelBuilder){
        String modelProvider = modelBuilder.getProvider();
        if (modelBuilder.wasDefault()){
            DEFAULT_MODEL_BUILDER_MAP.put(modelProvider,modelBuilder);
        }
        MODEL_BUILDER_MAP.add(modelProvider,modelBuilder);
    }

    public static ModelBuilder loadBuilder(AiNode aiNode, ModelParam params){
        List<ModelBuilder> list = MODEL_BUILDER_MAP.get(aiNode.getProvider());
        ModelBuilder targetBuilder = null;
        for (ModelBuilder ele : list){
            if (ele.checkExecModel(aiNode,params)){
                targetBuilder = ele;
                break;
            }
        }
        if (targetBuilder == null){
            targetBuilder = DEFAULT_MODEL_BUILDER_MAP.get(aiNode.getProvider());
        }
        return targetBuilder;
    }

    public static Collection supportProvider(){
        return MODEL_BUILDER_MAP.keySet();
    }

}
